from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from pydantic import BaseModel
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END

from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, MessagesState, StateGraph
from langchain_community.chat_models.tongyi import ChatTongyi
from langchain_core.messages import HumanMessage
import time
import chainlit as cl
from fastapi import FastAPI
from chainlit.utils import mount_chainlit
from chainlit.types import ThreadDict
from openai import AsyncOpenAI
from mcp import ClientSession
from typing import Dict, Optional
from fastapi import Request, Response
from chainlit.input_widget import Select, Switch, Slider
import pandas as pd
import plotly.graph_objects as go
import json
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver

 
model = ChatTongyi(
    model="qwen-max",   # 此处以qwen-max为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    streaming=True,
     api_key='sk-13c8bc2d23274db682f193b16ce57b64'
)

def tool() -> None:
    """Testing tool."""
    ...

def pre_model_hook() -> None:
    """Pre-model hook."""
    pass


def post_model_hook() -> None:
    """Post-model hook."""
    pass

class ResponseFormat(BaseModel):
    """Response format for the agent."""
    result: str

agent = create_react_agent(
    model,
    tools=[tool],
   # pre_model_hook=pre_model_hook,
   # post_model_hook=post_model_hook,
    response_format=ResponseFormat,
)

# Visualize the graph
# For Jupyter or GUI environments:
agent.get_graph().draw_mermaid_png()

# To save PNG to file:
png_data = agent.get_graph().draw_mermaid_png()
with open("graph.png", "wb") as f:
    f.write(png_data)

# For terminal/ASCII output:
agent.get_graph().draw_ascii()

for chunk in agent.stream(
    {"messages": [{"role": "user", "content": " 南京"}]},
    stream_mode="updates"
):
    print(chunk)